package com.navinfo.platform.wbs.service;

import com.mongodb.spark.MongoSpark;
import com.navinfo.platform.common.service.ICommonService;
import com.navinfo.platform.common.service.IDataHandleService;
import com.navinfo.platform.wbs.dto.OperationStatusData;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.MapFunction;
import org.apache.spark.api.java.function.MapPartitionsFunction;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import static org.apache.spark.sql.functions.col;
import static org.apache.spark.sql.functions.sum;

@SuppressWarnings("unchecked")
public enum OperationStatusMonthService implements IDataHandleService {
    INSTANCE;
    @Override
    public void handle(List<Object> obj, Map<String, String> configMap, SparkSession spark) {
        Dataset<Row> operationStatusDataset = (Dataset<Row>)obj.get(0);
        Dataset<Row> tid375Dataset = (Dataset<Row>)obj.get(1);
        Dataset<Row> tid500Dataset = (Dataset<Row>)obj.get(2);


        Dataset<Row> allDataset = operationStatusDataset.mapPartitions(getRowOperationStatusDataMapPartitionsFunction(), Encoders.bean(OperationStatusData.class))
        .groupBy("tid", "status")
        .agg(sum("duration").as("duration"));
        Dataset<OperationStatusData> tid375Data = allDataset
                .join(tid375Dataset, col("_c0").equalTo(col("tid")))
                .map(operationStatusDataMapFunction(375),
                        Encoders.bean(OperationStatusData.class));
        Dataset<Row> monthAgg375Data = tid375Data.groupBy("hpType", "status")
                .agg(sum("duration").as("duration"));

        Dataset<OperationStatusData> tid500Data = allDataset
                .join(tid500Dataset, col("_c0").equalTo(col("tid")))
                .map(operationStatusDataMapFunction(500),
                        Encoders.bean(OperationStatusData.class));
        Dataset<Row> monthAgg500Data = tid500Data.groupBy("hpType", "status")
                .agg(sum("duration").as("duration"));
        Dataset<Row> monthAggData = monthAgg375Data.union(monthAgg500Data);
        saveDataToMongo(configMap, monthAggData);
    }

    private MapPartitionsFunction<Row, OperationStatusData> getRowOperationStatusDataMapPartitionsFunction() {
        return f -> {
            List<OperationStatusData> operationStatusDataList = new ArrayList<>();
            while(f.hasNext()){
                Row row = f.next();
                OperationStatusData statusData = dataBeanprovider().call(row);
                int status = statusData.getStatus();
                if(status== 4){
                    long duration = statusData.getDuration();
                    if(duration <= 3600){
                        operationStatusDataList.add(statusData);
                    }
                } else {
                    operationStatusDataList.add(statusData);
                }
            }
            return operationStatusDataList.iterator();
        };
    }

    private static Function<Row, OperationStatusData> dataBeanprovider(){
        return row -> {
            OperationStatusData statusData = new OperationStatusData();
            statusData.setTid(row.getAs("tid"));
            statusData.setStatus(row.getAs("status"));
            int duration = row.getAs("duration");
            statusData.setDuration(duration);
            return statusData;
        };
    }

    private MapFunction<Row, OperationStatusData> operationStatusDataMapFunction(final int type) {
        return row ->{
            OperationStatusData operationStatusData = new OperationStatusData();
            operationStatusData.setTid(row.getAs("tid"));
            operationStatusData.setHpType(type);
            operationStatusData.setStatus(row.getAs("status"));
            operationStatusData.setDuration(row.getAs("duration"));
            return operationStatusData;
        };
    }

    private void saveDataToMongo(Map<String, String> config, Dataset<Row> saveData) {
        String day = config.get(ICommonService.DAY_VALUE);
        String collectionNamePrefix = config.get("mongo.collection.month.name");
        String collectionName = collectionNamePrefix + "_" + day.substring(2, 6);
        MongoSpark.write(saveData).option("collection", collectionName).mode("overwrite").save();
    }
}
